/***************************************************************
 * File : My_TCP.c
 * Description: lwIP-based TCP client for receiving a bin file
 *              from the Python server. Robust stream parsing,
 *              safe little-endian reads, ACK/NACK sending,
 *              and basic tcp_sent / tcp_err handling.
 *
 * 注意：
 *  - 需要在工程中链接 lwIP 并包含 lwIP 的头文件搜索路径
 *  - 根据目标平台和编译器，可能需要调整 printf 重定向
 ***************************************************************/
#include "My_TCP.h"  /* 如果你有额外头文件, 在此包含；否则把声明放到 .h 中 */
#include <lwip/tcp.h>
#include <lwip/ip_addr.h>
#include <lwip/memp.h>
#include <lwip/pbuf.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

#include "My_InFlash.h"
#include "crc32_table.h"
/* 定义端口号（可根据需要改）*/
#define TCP_LOCAL_PORT     2980
#define BLOCK_BUF_SIZE     (4096 + 16)
#define FILE_HEADER_LEN    20

/* 传输阶段枚举 */
typedef enum
{
    TCP_STAGE_WAIT_HEADER,
    TCP_STAGE_WAIT_BLOCK,
    TCP_STAGE_FINISH
} TcpTransferStage;

/* 接收状态结构体 */
typedef struct
{
    TcpTransferStage stage;
    uint8_t recv_buf[BLOCK_BUF_SIZE];
    uint32_t recv_len;
    uint32_t total_file_size;
    uint32_t total_blocks;
    uint32_t block_size;
    uint32_t global_crc32; // 服务器传来的整个文件的CRC32
    uint32_t current_block_idx;
    uint32_t accumulated_crc32; //本地累积计算的文件CRC32
} TcpRecvState;

static update_callback_t update_Complete_callback = NULL; //升级完成回调
void set_update_complete_callback(update_callback_t callback) {
    update_Complete_callback = callback;
}

static update_callback_t update_Begin_callback = NULL; //升级开始回调
void set_update_Begin_callback(update_callback_t callback) {
    update_Begin_callback = callback;
}

static update_callback_t update_Block_callback = NULL; //升级块回调
void set_update_Block_callback(update_callback_t callback) {
    update_Block_callback = callback;
}

/* ---------------- 辅助：小端安全读取 ---------------- */
static inline uint32_t read_u32_le(const uint8_t* p) {
    return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

static inline uint16_t read_u16_le(const uint8_t* p) {
    return (uint16_t)p[0] | ((uint16_t)p[1] << 8);
}

/* CRC16 (XMODEM)，与 Python 端匹配 */
static uint16_t crc16_xmodem(const uint8_t* data, uint32_t len) {
    uint16_t crc = 0;
    for (uint32_t i = 0; i < len; i++) {
        crc ^= (uint16_t)(data[i] << 8);
        for (uint8_t j = 0; j < 8; j++) {
            if (crc & 0x8000)
                crc = (crc << 1) ^ 0x1021;
            else
                crc <<= 1;
            crc &= 0xFFFF;
        }
    }
    return crc;
}


static uint32_t crc32_incremental(const uint8_t* data, uint32_t len, uint32_t prev_crc) {
    uint32_t crc = prev_crc ^ 0xFFFFFFFF;
    for (uint32_t i = 0; i < len; i++) {
        crc = (crc >> 8) ^ crc32_table[(crc ^ data[i]) & 0xFF];
    }
    return crc ^ 0xFFFFFFFF;
}

/* ---------------- 发送响应（ACK/NACK） ----------------
   格式：1 字节类型（0x00=ACK,0x01=NACK） + 4 字节块序号（小端）
   采用按字节填充，避免未对齐的 32 位写入问题
*/
static err_t send_response(struct tcp_pcb* tpcb, uint8_t resp_type, uint32_t block_idx) {
    uint8_t resp[5];
    resp[0] = resp_type;
    resp[1] = (uint8_t)(block_idx & 0xFF);
    resp[2] = (uint8_t)((block_idx >> 8) & 0xFF);
    resp[3] = (uint8_t)((block_idx >> 16) & 0xFF);
    resp[4] = (uint8_t)((block_idx >> 24) & 0xFF);

    err_t wr = tcp_write(tpcb, resp, sizeof(resp), TCP_WRITE_FLAG_COPY);
    if (wr != ERR_OK) {
        /* 常见原因：ERR_MEM（发送队列已满）等 */
        printf("tcp_write 响应失败, err=%d\r\n", wr);
        /* 尝试触发底层发送，以便释放发送队列（视 lwIP 配置而定） */
        tcp_output(tpcb);
        return wr;
    }
    tcp_output(tpcb);
    printf("发送响应：type=%u, idx=%u\r\n", resp_type, block_idx);
    return ERR_OK;
}

/* ---------------- tcp_sent 回调（当发送队列已发送出去时触发） -----------
   可用于实现更复杂的重发队列/释放缓冲的逻辑；
   此处仅打印日志，必要时你可以扩展为重试队列。
*/
static err_t tcp_client_sent(void* arg, struct tcp_pcb* tpcb, u16_t len) {
    (void)arg;
    (void)len;
    /* 这里 len 表示 lwIP 已确认发送出去了多少字节 */
    // printf("tcp_sent 回调：%u 字节已发送\r\n", (unsigned)len);
    return ERR_OK;
}

/* ---------------- tcp_err 回调（连接被底层关闭或发生错误） -----------
   在 err 回调中不能使用 tpcb（已经被释放），只释放分配的 state。
*/
static void tcp_client_err(void* arg, err_t err) {
    TcpRecvState* state = (TcpRecvState*)arg;
    printf("tcp_err 回调，err=%d，释放状态。\r\n", err);
    if (state) {
        free(state);
    }
}

/* ---------------- 核心接收回调（流式解析，支持多块同时到达） -------------- */
static err_t tcp_client_recv(void* arg, struct tcp_pcb* tpcb, struct pbuf* p, err_t err) {
    TcpRecvState* state = (TcpRecvState*)arg;

    if (p == NULL) {
        /* 对端主动关闭连接（FIN），释放资源 */
        printf("TCP 连接关闭通知\r\n");
        if (state)
            free(state);
        return tcp_close(tpcb);
    }

    if (err != ERR_OK) {
        printf("接收回调 err != ERR_OK : %d\r\n", err);
        pbuf_free(p);
        return err;
    }

    /* 1) 将 pbuf 链数据安全拷贝到 state->recv_buf（避免越界） */
    uint32_t copied = 0;
    struct pbuf* q = p;
    while (q != NULL) {
        uint32_t remain = BLOCK_BUF_SIZE - state->recv_len;
        if (remain == 0) {
            /* 缓冲区满了 —— 这是异常情况（数据超出了缓冲），直接跳出并丢弃后续 */
            printf("接收缓冲区已满，丢弃后续 %u 字节\r\n", (unsigned)q->len);
            break;
        }
        uint32_t to_copy = (q->len > remain) ? remain : q->len;
        memcpy(state->recv_buf + state->recv_len, q->payload, to_copy);
        state->recv_len += to_copy;
        copied += to_copy;
        q = q->next;
    }
    printf("接收数据: %u bytes (累计: %u)\r\n", (unsigned)copied, (unsigned)state->recv_len);

    /* 2) 循环解析流内可能的多条消息（header 或多个块） */
    while (1) {
        if (state->stage == TCP_STAGE_WAIT_HEADER) {
            if (state->recv_len >= FILE_HEADER_LEN) {
                uint32_t magic = read_u32_le(state->recv_buf);
                if (magic != 0x55AA55AAU) {
                    printf("文件头错误，魔数不匹配（实际 0x%08X）\r\n", magic);
                    pbuf_free(p);
                    return ERR_VAL;
                }

                state->total_file_size = read_u32_le(state->recv_buf + 4);
                state->total_blocks = read_u32_le(state->recv_buf + 8);
                state->block_size = read_u32_le(state->recv_buf + 12);
                state->global_crc32 = read_u32_le(state->recv_buf + 16);
                state->accumulated_crc32 = 0;
                printf("文件头解析成功：size=%u, blocks=%u, block_size=%u, CRC32=0x%08X\r\n",
                       state->total_file_size, state->total_blocks, state->block_size, state->global_crc32);
                printf("收到升级请求成功，切换为接受数据块...并且擦除指定地址\r\n");
                if (update_Begin_callback) {
                    update_Begin_callback(1);
                }

                if (InternalFlash_Erase(APP2, 256 * 1024, 2)) {
                    printf("擦除失败\r\n");
                }


                state->stage = TCP_STAGE_WAIT_BLOCK;
                state->current_block_idx = 0;

                /* 消耗 header 字节，保留剩余 */
                uint32_t remain = state->recv_len - FILE_HEADER_LEN;
                if (remain > 0) {
                    memmove(state->recv_buf, state->recv_buf + FILE_HEADER_LEN, remain);
                }
                state->recv_len = remain;
                /* 继续循环解析（若缓冲中已有块） */
                continue;
            } else {
                /* header 未到齐，等待下一次数据 */
                break;
            }
        } else if (state->stage == TCP_STAGE_WAIT_BLOCK) {
            if (state->recv_len < 6) {
                /* 甚至连块头都不完整，等待更多数据 */
                break;
            }
            uint32_t block_idx = read_u32_le(state->recv_buf);
            uint16_t data_len = read_u16_le(state->recv_buf + 4);
            uint32_t total_block_len = 6 + (uint32_t)data_len + 2; /* 头 + 数据 + CRC16 */

            if (state->recv_len < total_block_len) {
                /* 当前块不完整，等待下一次 recv */
                break;
            }

            /* 完整块到达 */
            uint8_t* block_data = state->recv_buf + 6;
            uint16_t recv_crc16 = read_u16_le(state->recv_buf + 6 + data_len);

            uint8_t resp_type = 0x01; /* NACK 默认 */
            if (block_idx == state->current_block_idx) {
                uint16_t calc_crc = crc16_xmodem(block_data, data_len);
                if (calc_crc == recv_crc16) {
                    resp_type = 0x00; /* ACK */
                    printf("块 %u 校验成功 (len=%u CRC=0x%04X)\r\n", block_idx, data_len, calc_crc);
                    state->accumulated_crc32 = crc32_incremental(block_data, data_len, state->accumulated_crc32);
                    printf("块 %u 累积CRC32: 0x%08X\r\n", block_idx, state->accumulated_crc32);
                    if (update_Block_callback) {
                        update_Block_callback(1);
                    }

                    // 写入 Flash 或文件：
                    uint8_t err = InFlash_Write(
                        APP2 + (block_idx * state->block_size), block_data, data_len);
                    if (err) {
                        printf("InFlash_Write err : %d", err);
                    }

                    state->current_block_idx++;
                } else {
                    printf("块 %u CRC 错误，接收 0x%04X，计算 0x%04X\r\n", block_idx, recv_crc16, calc_crc);
                }
            } else {
                printf("块序号错误：接收 %u，期望 %u\r\n", block_idx, state->current_block_idx);
            }

            /* 发送 ACK/NACK（如果发送失败，返回值会指出原因） */
            err_t sres = send_response(tpcb, resp_type, block_idx);
            if (sres != ERR_OK) {
                /* 简单记录。复杂场景可把未发送的响应入队，并在 tcp_sent 回调中重试 */
                printf("发送响应返回 err=%d\r\n", sres);
            }

            /* 消耗已处理的字节，并把剩余数据 memmove 到缓冲头 */
            uint32_t remain_after = state->recv_len - total_block_len;
            if (remain_after > 0) {
                memmove(state->recv_buf, state->recv_buf + total_block_len, remain_after);
            }
            state->recv_len = remain_after;

            /* 检查是否全部块接收完成 */
            if (state->current_block_idx >= state->total_blocks) {
                printf("所有块接收完成，进入整体校验阶段\r\n");
                state->stage = TCP_STAGE_FINISH;

                // 比对本地计算的CRC32与服务器传来的global_crc32
                uint8_t final_resp = (state->accumulated_crc32 == state->global_crc32) ? 0x00 : 0x01;
                printf("整体CRC校验：本地=0x%08X，服务器=0x%08X，结果=%s\r\n",
                       state->accumulated_crc32, state->global_crc32,
                       final_resp == 0x00 ? "成功" : "失败");
                // 调用回调函数通知上层
                if (update_Complete_callback) {
                    update_Complete_callback(final_resp);
                }

                // 发送最终结果给服务器
                err_t wr = tcp_write(tpcb, &final_resp, 1, TCP_WRITE_FLAG_COPY);
                if (wr == ERR_OK)
                    tcp_output(tpcb);
                printf("已发送整体校验结果：%u (tcp_write err=%d)\r\n", final_resp, (int)wr);
            }
            /* 继续 while 循环以处理缓冲中可能的下一个块 */
            continue;
        } else {
            /* TCP_STAGE_FINISH 或其他，暂时不处理更多 */
            break;
        }
    }

    /* 告诉 lwIP 我们已接收 p->tot_len 字节，并释放 pbuf */
    tcp_recved(tpcb, p->tot_len);
    pbuf_free(p);
    return ERR_OK;
}

/* ---------------- 连接建立回调 ---------------- */
static uint8_t TCP_connected_Flag = 0;

static err_t tcp_client_connected(void* arg, struct tcp_pcb* tpcb, err_t err) {
    (void)arg;
    if (err != ERR_OK) {
        printf("连接失败，err=%d\r\n", err);
        return err;
    }

    TCP_connected_Flag = 1;
    printf("TCP 客户端已连接到服务器\r\n");

    /* 注册回调（recv, sent, err） */
    tcp_recv(tpcb, tcp_client_recv);
    tcp_sent(tpcb, tcp_client_sent);
    tcp_err(tpcb, tcp_client_err);

    return ERR_OK;
}


struct tcp_pcb* tcp_client_init(uint8_t Ip[4], uint16_t LOCAL_PORT, uint16_t REMOTE_PORT) {
    ip_addr_t serverIp;
    struct tcp_pcb* tpcb = NULL;
    TcpRecvState* state = NULL;

    /* 分配状态结构体 */
    state = (TcpRecvState*)malloc(sizeof(TcpRecvState));
    if (state == NULL) {
        printf("内存分配失败 (state)\r\n");
        return NULL;
    }
    memset(state, 0, sizeof(TcpRecvState));
    state->stage = TCP_STAGE_WAIT_HEADER;
    state->recv_len = 0;
    state->current_block_idx = 0;

    /* 创建 TCP 控制块 */
    tpcb = tcp_new();
    if (tpcb == NULL) {
        printf("tcp_new 失败\r\n");
        free(state);
        return NULL;
    }

    /* 将 state 绑定到 tpcb 以便回调使用 */
    tcp_arg(tpcb, state);

    /* 绑定本地端口（可选） */
    err_t bind_err = tcp_bind(tpcb, IP_ADDR_ANY, LOCAL_PORT);
    if (bind_err != ERR_OK) {
        printf("绑定本地端口失败 err=%d\r\n", bind_err);
        tcp_abort(tpcb);
        free(state);
        return NULL;
    }

    /* 连接到服务器 */
    IP4_ADDR(&serverIp, Ip[0], Ip[1], Ip[2], Ip[3]);
    err_t conn_err = tcp_connect(tpcb, &serverIp, REMOTE_PORT, tcp_client_connected);
    if (conn_err != ERR_OK) {
        printf("tcp_connect 失败 err=%d\r\n", conn_err);
        tcp_abort(tpcb);
        free(state);
        return NULL;
    }
    printf("TCP 客户端初始化完成，正在连接 %u.%u.%u.%u:%u ...\r\n", Ip[0], Ip[1], Ip[2], Ip[3], REMOTE_PORT);
    return tpcb;
}
